import logging
from typing import Any, Dict, List, Tuple
import pandas as pd
import aiohttp
import asyncio

# ===================== 全局配置 =====================
API_URL: str = "http://GenBiM-GenBi-c8nytehOGcCE-375297461.cn-northwest-1.elb.amazonaws.com.cn:8088/index/docs"
PROFILE_NAME: str = "经营分析"
INDEX_NAME: str = "ner_index"
EXCEL_PATH: str = "data.xlsx"
ENTITY_COL: str = "entity"
COMMENT_COL: str = "comment"
RESULT_COL: str = "result"
MAX_CONCURRENCY: int = 8
TIMEOUT: int = 10

# ===================== 日志配置 =====================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s"
)
logger = logging.getLogger(__name__)

# ===================== 单条异步请求函数 =====================
async def post_doc(session: aiohttp.ClientSession, question: str, answer: str) -> Tuple[bool, str]:
    """
    发起单条POST异步请求，返回(是否成功, 返回内容或错误信息)
    """
    payload: Dict[str, Any] = {
        "profile_name": PROFILE_NAME,
        "question": question,
        "answer": answer,
        "index_name": INDEX_NAME
    }
    headers: Dict[str, str] = {
        "Accept": "*/*",
        "Content-Type": "application/json"
    }
    try:
        async with session.post(API_URL, json=payload, headers=headers, timeout=TIMEOUT) as resp:
            text = await resp.text()
            if resp.status == 200:
                return True, text
            else:
                logger.error(f"请求失败: {resp.status} {text}")
                return False, f"{resp.status}: {text}"
    except Exception as e:
        logger.error(f"请求异常: {e}")
        print(f'请求异常：{e}')
        return False, str(e)

# ===================== 批量异步处理主流程 =====================
async def batch_process(df: pd.DataFrame) -> pd.DataFrame:
    """
    并发批量异步请求，将结果写入新列
    """
    results: List[str] = ["" for _ in range(len(df))]
    connector = aiohttp.TCPConnector(limit=MAX_CONCURRENCY)
    timeout = aiohttp.ClientTimeout(total=None, sock_connect=TIMEOUT, sock_read=TIMEOUT)
    async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
        tasks = [post_doc(session, str(row[ENTITY_COL]), str(row[COMMENT_COL])) for _, row in df.iterrows()]
        for idx, coro in enumerate(asyncio.as_completed(tasks)):
            try:
                success, msg = await coro
                results[idx] = f"success: {msg}" if success else f"fail: {msg}"
            except Exception as e:
                logger.error(f"异步线程异常: {e}")
                print(f"异步线程异常: {e}")
                results[idx] = f"fail: {e}"
    df[RESULT_COL] = results
    return df

# ===================== 主程序入口 =====================
async def async_main() -> None:
    """
    读取Excel，批量异步请求，写回结果
    """
    try:
        df = pd.read_excel(EXCEL_PATH)
    except Exception as e:
        logger.error(f"读取Excel失败: {e}")
        return
    if ENTITY_COL not in df.columns or COMMENT_COL not in df.columns:
        logger.error(f"Excel缺少必要列: {ENTITY_COL}, {COMMENT_COL}")
        return
    logger.info(f"读取到{len(df)}条数据，开始批量异步请求...")
    df = await batch_process(df)
    try:
        df.to_excel(EXCEL_PATH, index=False)
        logger.info(f"结果已写回 {EXCEL_PATH} 的 {RESULT_COL} 列")
    except Exception as e:
        logger.error(f"写回Excel失败: {e}")

# ===================== 使用示例 =====================
if __name__ == "__main__":
    asyncio.run(async_main())
